import os
import time
from datetime import timedelta
from typing import Dict, Any, Generator, List, Optional
import easyocr
from openai import OpenAI, APIError
from PyPDF2 import PdfReader
from docx import Document
import numpy as np
import cv2


def classify_contract_type(text: str) -> str:
    categories = {
        # 按期限
        "固定期限劳动合同": ["合同期限自", "为期", "终止时间"],
        "无固定期限劳动合同": ["无固定期限", "不设终止日期", "除非双方解除"],
        "以完成工作任务为期限劳动合同": ["任务完成之日", "项目结束即终止", "工作成果"],
        # 按工作时间强度
        "全日制劳动合同": ["标准工时制", "五天八小时", "每日工作时长"],
        "非全日制用工合同": ["非全日制", "小时工", "计时工", "累计工作时间"],
        "零工/临时工合同": ["临时", "一次性任务", "单次派工", "劳务报酬"],
        # 按用工主体
        "劳务派遣合同": ["派遣公司", "派遣单位", "劳务派遣", "第三方用工"],
        "劳务外包合同": ["承揽", "外包", "服务费"],
        # 按身份资历
        "实习/见习协议": ["实习", "见习", "学校三方协议", "实习期间"],
        "试用期协议": ["试用期", "不符合录用条件", "提前通知解除"],
        "顾问/专家聘用合同": ["顾问", "专家", "咨询服务", "顾问费"],
        "兼职协议": ["兼职", "兼任", "工作时间安排"],
        # 专项条款（独立或附属）
        "保密协议（NDA）": ["保密协议", "保密信息", "违约责任", "保密期限"],
        "竞业限制协议": ["竞业限制", "竞业限制期", "经济补偿", "地域范围"],
        "员工持股/股权激励协议": ["股权激励", "员工持股", "认购价格", "归属期限"],
        "劳动争议调解/赔偿协议": ["争议解决", "和解", "一次性赔偿", "调解"],
        # 新兴用工模式
        "平台经济用工协议": ["平台", "派单", "灵活用工", "平台与劳动者"],
        "远程/异地用工合同": ["远程办公", "异地", "通信工具", "工作地点"],
        "弹性用工协议": ["弹性工时", "核心工作时段", "自主排班"],
    }

    for name, keywords in categories.items():
        if any(keyword in text for keyword in keywords):
            return name
    return "无法识别/可能为非标准劳动合同"


def detect_seal_in_text(text: str) -> bool:
    """通过关键字判断 OCR 文本中是否有‘公章’、‘盖章’等字样。"""
    seal_keywords = ["公章", "盖章", "（章）", "企业章", "法人章"]
    return any(k in text for k in seal_keywords)

def detect_seal_in_image(file_path: str, min_area_ratio: float = 0.001) -> bool:
    """
    对单张图片做红色圆形印章的更严格检测：
    - 颜色分割后要求红色区域占比超过 min_area_ratio
    - Hough 圆检测验证圆形结构
    """
    if not os.path.exists(file_path):
        return False
    arr = np.fromfile(file_path, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        return False
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    # 红色在 HSV 空间范围
    lower1, upper1 = np.array([0, 70, 50]), np.array([10, 255, 255])
    lower2, upper2 = np.array([170, 70, 50]), np.array([180, 255, 255])
    mask = cv2.inRange(hsv, lower1, upper1) | cv2.inRange(hsv, lower2, upper2)
    # 计算红色像素比例
    red_ratio = np.sum(mask > 0) / (mask.size + 1e-6)
    if red_ratio < min_area_ratio:
        return False
    # 模糊后找圆
    blurred = cv2.GaussianBlur(mask, (9, 9), 2)
    circles = cv2.HoughCircles(
        blurred, cv2.HOUGH_GRADIENT, dp=1.2, minDist=100,
        param1=50, param2=35, minRadius=30, maxRadius=150
    )
    return circles is not None


class ContractAnalyzer:
    def __init__(self, api_key: str):
        self.client = OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )

    def _extract_from_image(self, file_path: str, reader=None) -> str:
        try:
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"文件不存在: {file_path}")
            img_array = np.fromfile(file_path, dtype=np.uint8)
            image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
            if image is None:
                raise ValueError("无法解析图像，cv2.imdecode 失败")
            reader = reader or easyocr.Reader(['ch_sim', 'en'])
            result = reader.readtext(image, detail=0)
            return "\n".join(result).strip()
        except Exception as e:
            raise RuntimeError(f"图片OCR识别失败: {str(e)}")

    def extract_text(self, file_path: str) -> str:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")
        ext = file_path.lower()
        if ext.endswith('.pdf'):
            return self._extract_from_pdf(file_path)
        elif ext.endswith('.docx'):
            return self._extract_from_docx(file_path)
        elif ext.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp')):
            return self._extract_from_image(file_path)
        else:
            raise ValueError("不支持的文件格式，请提供PDF、DOCX或图片文件")

    def extract_texts_from_multiple_images(self, image_paths: List[str]) -> str:
        reader = easyocr.Reader(['ch_sim', 'en'])
        all_text = []
        for path in image_paths:
            try:
                text = self._extract_from_image(path, reader)
                all_text.append(text)
            except Exception as e:
                print(f"\n⚠️ 图片处理失败 [{path}]: {str(e)}")
        return "\n\n".join(all_text)

    def _extract_from_pdf(self, file_path: str) -> str:
        text = ""
        try:
            with open(file_path, 'rb') as file:
                reader = PdfReader(file)
                for page in reader.pages:
                    text += page.extract_text() or ""
        except Exception as e:
            raise RuntimeError(f"PDF文件读取失败: {str(e)}")
        return text

    def _extract_from_docx(self, file_path: str) -> str:
        try:
            doc = Document(file_path)
            return "\n".join([para.text for para in doc.paragraphs if para.text])
        except Exception as e:
            raise RuntimeError(f"DOCX文件读取失败: {str(e)}")

    def analyze_contract_stream(
        self,
        text: str,
        contract_type: str,
        file_paths: Optional[List[str]] = None
    ) -> Generator[str, None, Dict[str, Any]]:
        if not text.strip():
            raise ValueError("合同文本内容为空")

        # 文本层面印章检测
        has_seal_text = detect_seal_in_text(text)
        # 图像层面印章检测
        has_seal_image = False
        if file_paths:
            for p in file_paths:
                if p.lower().endswith(('.png','.jpg','.jpeg','.bmp','.tiff')) and detect_seal_in_image(p):
                    has_seal_image = True
                    break

        seal_note = "（检测到公章）" if has_seal_text and has_seal_image else "（未检测到公章）"
        print(seal_note)

        system_prompt = f"""你是一位专业的劳动法律师，负责分析劳动合同的合规性。
该合同初步识别为：{contract_type} {seal_note}

请严格按照以下要求进行分析：
1. 首先判断并确认合同类型（是否准确）
2. 给出总体评价（正规/基本正规/存在明显问题）
3. 根据 {seal_note}说明是否有公章
4. 然后分条列出具体问题或漏洞
5. 对每个问题，简要说明法律依据或建议
6. 重点检查以下关键条款：
   - 合同双方基本信息是否完整
   - 劳动合同期限
   - 工作内容和工作地点
   - 工作时间和休息休假
   - 劳动报酬
   - 社会保险
   - 劳动保护、劳动条件和职业危害防护
   - 劳动合同解除或终止条件
   - 违约责任
   - 竞业限制条款（如有）
   - 保密条款（如有）
"""

        try:
            stream = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": f"请分析以下劳动合同：\n{text[:15000]}"}
                ],
                temperature=0.3,
                max_tokens=2000,
                stream=True
            )

            collected_content = []
            total_tokens = 0
            for chunk in stream:
                if not chunk.choices:
                    continue
                delta = chunk.choices[0].delta
                if delta and delta.content:
                    collected_content.append(delta.content)
                    yield delta.content
                if hasattr(chunk, 'usage') and chunk.usage:
                    total_tokens = chunk.usage.total_tokens

            return {"metadata": {"total_tokens": total_tokens, "complete_response": "".join(collected_content)}}

        except APIError as e:
            raise RuntimeError(f"API请求失败: {str(e)}")
        except Exception as e:
            raise RuntimeError(f"分析过程中出错: {str(e)}")


def main():
    API_KEY = "sk-20856422ed6644e3827b9d5403c9542a"  # 替换为你的API密钥
    analyzer = ContractAnalyzer(API_KEY)

    print("劳动合同分析工具（流式输出版）")
    print("=" * 40)
    file_input = input("请输入劳动合同路径（多个图片用英文逗号分隔，或一个PDF/DOCX）: ").strip()
    file_paths = [p.strip() for p in file_input.split(',') if p.strip()]

    if not file_paths:
        print("❌ 未输入有效路径")
        return

    # 判断模式
    if len(file_paths) > 1:
        if all(p.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.tiff')) for p in file_paths):
            mode = "multi_image"
        else:
            print("❌ 当前仅支持：多张图片 或 单个 PDF/DOCX。请检查输入。")
            return
    elif file_paths[0].lower().endswith(('.pdf', '.docx')):
        mode = "single_document"
    else:
        print("❌ 当前仅支持：多张图片 或 单个 PDF/DOCX。请检查输入。")
        return

    try:
        # 文本提取
        if mode == "multi_image":
            text = analyzer.extract_texts_from_multiple_images(file_paths)
        else:
            text = analyzer.extract_text(file_paths[0])

        if not text.strip():
            print("❌ 没有提取到有效文本，终止分析")
            return

        # 初步识别类型
        contract_type = classify_contract_type(text)
        print(f"\n📌 初步识别的合同类型：{contract_type}")

        print("\n正在分析合同，请稍候...\n")
        print("=" * 40)
        print("实时分析结果:")
        print("=" * 40)

        start_time = time.time()
        full_response = []
        metadata = {}

        # 分析并输出
        for chunk in analyzer.analyze_contract_stream(text, contract_type, file_paths):
            print(chunk, end="", flush=True)
            full_response.append(chunk)
        elapsed = time.time() - start_time

        print("\n\n" + "=" * 40)
        print("分析完成!")
        if metadata and metadata.get('total_tokens', 0) > 0:
            print(f"\n总Tokens: {metadata['total_tokens']}")
            print(f"总耗时: {timedelta(seconds=elapsed)}")
            print(f"处理速度: {metadata['total_tokens']/elapsed:.2f} tokens/秒")

        # 保存结果
        save_path = os.path.join(os.path.dirname(file_paths[0]), "analysis_result_stream.txt")
        with open(save_path, 'w', encoding='utf-8') as f:
            f.write("".join(full_response))
        print(f"\n✅ 分析结果已保存到: {save_path}")

    except KeyboardInterrupt:
        print("\n用户中断操作")
    except Exception as e:
        print(f"\n程序发生错误: {str(e)}")


if __name__ == "__main__":
    main()
